package org.mmtk.policy.compressor;

import static org.mmtk.policy.compressor.Chunk.LOG_BYTES_IN_CHUNK;
import static org.mmtk.policy.compressor.Block.LOG_BYTES_IN_BLOCK;

import org.mmtk.policy.Space;
import org.mmtk.utility.Constants;
import org.mmtk.utility.alloc.EmbeddedMetaData;

import org.mmtk.vm.VM;
import org.vmmagic.pragma.Inline;
import org.vmmagic.pragma.Uninterruptible;
import org.vmmagic.unboxed.Address;
import org.vmmagic.unboxed.Extent;
import org.vmmagic.unboxed.ObjectReference;
import org.vmmagic.unboxed.Word;

@Uninterruptible
public class CompressorMetadata implements Constants {
  
  //
  // "Data Metadata" is the metadata of the Compressor. It's memory map looks like this
  // +-------------------------------------+----------------------+--------------+------------+
  // | LiveBits Map                        | Offset Vector        | FirstObj Vec | Status Vec |
  // +-------------------------------------+----------------------+--------------+------------+
  //
  // The general memory map looks like this:
  // Note: Region == Block != Chunk (4MB)
  //
  //         +-------------+-------------+-------------+-----------------
  //         | Region  End | Next Region |  Data  End  | 
  //         +-------------+-------------+-------------+- Data --> 
  //         | Region Metadata (BumpPointer metadata)  | 
  //         +-----------------------------------------+----------------
  //         |                                                       /
  //         |                +-------------------------------------+
  //         /               /
  //  +-----+---------------+---------------------------
  //  |     | Block1        |  Block2   ...
  //  +-----+------------------------------------------
  //  | (*) |  Chunk data (BumpPointer Blocks)
  //  +-----+------------------------------------------
  //  (*) Chunk metadata
  

  
  // LOG_MIN_ALIGNMENT = 2 --> 4 byte resolution
  private static final int OBJECT_LIVE_SHIFT = LOG_MIN_ALIGNMENT; 
  private static final int LOG_BIT_COVERAGE = OBJECT_LIVE_SHIFT;

  // A value of 5 would indicate a coverage of 32 
  // 1 metadata unit for every 32 unit of data. Units are decided by alignment
  public static final int LOG_LIVE_COVERAGE = LOG_BIT_COVERAGE + LOG_BITS_IN_BYTE;
  public static final int LOG_OFFSET_COVERAGE = LOG_BYTES_IN_BLOCK - LOG_BYTES_IN_ADDRESS;
  
  // FIRSTOBJ_COVERAGE = BYTES_IN_PAGE / BYTES_IN_ADDRESS
  // 1 byte of metadata for each FIRSTOBJ_COVERAGE bytes
  // so BYTES_IN_ADDRESS bytes of metadata for each BYTES_IN_ADDRESS * FIRSTOBJ_COVERAGE bytes
  public static final int LOG_FIRSTOBJ_COVERAGE = LOG_BYTES_IN_PAGE - LOG_BYTES_IN_ADDRESS;
  public static final int LOG_STATUS_COVERAGE = LOG_BYTES_IN_CHUNK - LOG_BYTES_IN_ADDRESS;
  
  // How many live bytes there are per-chunk
  public static final int LOG_LIVE_BYTES_PER_CHUNK = LOG_BYTES_IN_CHUNK - LOG_LIVE_COVERAGE;
  public static final int LIVE_BYTES_PER_CHUNK = 1 << LOG_LIVE_BYTES_PER_CHUNK;
  public static final Extent LIVE_BYTES_EXTENT = Extent.fromIntSignExtend(LIVE_BYTES_PER_CHUNK);
  public static final Extent LIVE_BYTES_OFFSET = Extent.zero();
  
  // An address size (pointer) for each Block in the chunk
  public static final int OFFSET_VEC_BYTES_PER_CHUNK = 1 << (LOG_BYTES_IN_CHUNK - LOG_OFFSET_COVERAGE);
  public static final Extent OFFSET_VEC_BYTES_EXTENT = Extent.fromIntSignExtend(OFFSET_VEC_BYTES_PER_CHUNK);
  public static final Extent OFFSET_VEC_OFFSET = LIVE_BYTES_EXTENT;
  
  // An address size (pointer) for each Page in a chunk
  public static final int FIRSTOBJ_VEC_BYTES_PER_CHUNK = 1 << (LOG_BYTES_IN_CHUNK - LOG_FIRSTOBJ_COVERAGE);
  public static final Extent FIRSTOBJ_VEC_BYTES_EXTENT = Extent.fromIntSignExtend(FIRSTOBJ_VEC_BYTES_PER_CHUNK);
  public static final Extent FIRSTOBJ_VEC_OFFSET = OFFSET_VEC_OFFSET.plus(OFFSET_VEC_BYTES_EXTENT);
  
  // A byte for each Page in the chunk to hold status (one of HANDLED, UNHANDLED, BUSY)
  public static final int STATUS_VEC_BYTES_PER_CHUNK = 1 << (LOG_BYTES_IN_CHUNK - LOG_STATUS_COVERAGE);
  public static final Extent STATUS_VEC_BYTES_EXTENT = Extent.fromIntSignExtend(STATUS_VEC_BYTES_PER_CHUNK);
  public static final Extent STATUS_VEC_OFFSET = FIRSTOBJ_VEC_OFFSET.plus(FIRSTOBJ_VEC_BYTES_EXTENT);
  
  public static final int TOTAL_METADATA_BYTES_PER_CHUNK = LIVE_BYTES_PER_CHUNK + OFFSET_VEC_BYTES_PER_CHUNK + 
  FIRSTOBJ_VEC_BYTES_PER_CHUNK + STATUS_VEC_BYTES_PER_CHUNK;
  
  // (1 << LOG_BITS_IN_WORD) - 1, e.g. 0x1F if LOG_BITS_IN_WORD = 5
  private static final Word WORD_SHIFT_MASK = Word.one().lsh(LOG_BITS_IN_WORD).minus(Extent.one());
  
  /**
   * Atomically set the live bits for a given object at the start and the end of the object
   * This action is linearized wrt the second live bit set.
   *
   * @param object The object whose live bit is to be set.
   * @return True if the second bit was changed to true.
   */
  @Inline
  public static boolean testAndSetLiveBits(ObjectReference object) {
    // TODO Compressor article states that we should mark both the start and the end of an object,
    // however currently marking only the start is sufficient 
    
    Address objStart = VM.objectModel.objectStartRef(object);
    // objEnd = VM.objectModel.getObjectEndAddress(object).minus(1);
    //Address objEnd = objStart.plus(VM.objectModel.getCurrentSize(object));
    //updateLiveBit(objEnd, true, true);
    return updateLiveBit(objStart, true, true);
  }
  
  @Inline
  public static boolean testAndClearLiveBits(ObjectReference object) {
    // TODO Compressor article states that we should mark both the start and the end of an object,
    // however currently marking only the start is sufficient 
    
    Address objStart = VM.objectModel.objectStartRef(object);
    // objEnd = VM.objectModel.getObjectEndAddress(object).minus(1);
    //Address objEnd = objStart.plus(VM.objectModel.getCurrentSize(object));
    //updateLiveBit(objEnd, true, true);
    return updateLiveBit(objStart, false, true);
  }
  
  /**
   * Set the live bit for a given address
   *
   * @param address The address whose live bit is to be set.
   * @param set True if the bit is to be set, as opposed to cleared
   * @param atomic True if we want to perform this operation atomically
   * 
   * @return true iff value was changed, false otherwise
   */
  @Inline
  private static boolean updateLiveBit(Address address, boolean set, boolean atomic) {
    Word oldValue, newValue;
    Address liveWord = getLiveWordAddress(address);
    Word mask = getMask(address, true);
    if (atomic) {
      do {
        oldValue = liveWord.prepareWord();
        newValue = (set) ? oldValue.or(mask) : oldValue.and(mask.not());
      } while (!liveWord.attempt(oldValue, newValue));
    } else {
      oldValue = liveWord.loadWord();
      liveWord.store(set ? oldValue.or(mask) : oldValue.and(mask.not()));
    }
    return oldValue.and(mask).NE(mask);
  }


  /**
   * Test the live bit for a given object
   *
   * @param object The object whose live bit is to be checked
   * @return true iff the live bit is set
   */
  @Inline
  protected static boolean liveBitSet(ObjectReference object) {
    return liveBitSet(VM.objectModel.objectStartRef(object));
  }

  /**
   * Test the live bit for a given address
   *
   * @param object The object whose live bit is to be checked
   * @return true iff the live bit is set
   */
  @Inline
  protected static boolean liveBitSet(Address address) {
    Address liveWord = getLiveWordAddress(address);
    Word mask = getMask(address, true);
    Word value = liveWord.loadWord();
    return value.and(mask).EQ(mask);
  }

  /**
   * Clear the live bit for a given object
   *
   * @param object The object whose live bit is to be cleared.
   */
  @Inline
  public static void clearLiveBits(ObjectReference object) {
    // TODO Compressor article states that we should mark both the start and the end of an object,
    // however currently marking only the start is sufficient 
    Address objStart = VM.objectModel.objectStartRef(object);
    //Address objEnd = VM.objectModel.getObjectEndAddress(object).minus(1);
    clearLiveBit(objStart);
    //clearLiveBit(objEnd);
  }

  /**
   * Clear the live bit for a given address
   *
   * @param address The address whose live bit is to be cleared.
   */
  @Inline
  private static void clearLiveBit(Address address) {
    updateLiveBit(address, false, true);
  }
  
  /**
   * Given an address, produce a bit mask for the live table
   *
   * @param address The address whose live bit mask is to be established
   * @param set True if we want the mask for <i>setting</i> the bit,
   * false if we want the mask for <i>clearing</i> the bit.
   * @return The appropriate bit mask for object for the live table for.
   */
  @Inline
  public static Word getMask(Address address, boolean set) {
    //OBJECT_LIVE_SHIFT = 2
    //WORD_SHIFT_MASK = 0x1f = (5 lsb bits on)
    // Removes the 2 lsb and takes 5 bits (which index 32 bytes)
    // 5 bits represent a number between 0-31 which is the index of the bit
    // whose mask we are looking for -> shift is that index
    int shift = address.toWord().rshl(OBJECT_LIVE_SHIFT).and(WORD_SHIFT_MASK).toInt();
    Word rtn = Word.one().lsh(shift);
    return (set) ? rtn : rtn.not();
  }
  
  /**
   *                                            
   * @param address address the address to get the base for
   * @return The base of the data metadata for the chunk corresponding to address
   */
  public static Address getMetadataBase(Address address) {
    if (VM.VERIFY_ASSERTIONS) {
      VM.assertions._assert(!address.isZero());
      VM.assertions._assert(Space.chunkAlign(address, true).EQ(EmbeddedMetaData.getMetaDataBase(address)));
      VM.assertions._assert(EmbeddedMetaData.getMetaDataBase(address).EQ(Chunk.align(address)));
    }
    
    return Chunk.align(address);
  }
  
  /**
   * Returns the start address of the data section of the chunk where address resides (first address after the end of chunk metadata)
   * @param address an address
   * @return the start address of the data section of the chunk where address resides (first address after the end of chunk metadata)
   */
  public static Address getDataStart(Address address) {
    if (VM.VERIFY_ASSERTIONS) {
      VM.assertions._assert(!address.isZero());
    }
    
    return getMetadataBase(address).plus(TOTAL_METADATA_BYTES_PER_CHUNK);
  }
  
  public static Address getLiveMapBase(Address address) {
    return getMetadataBase(address).plus(LIVE_BYTES_OFFSET);
  }
  
  public static Address getOffsetVecBase(Address address) {
    // Comes exactly after liveMap
    return getMetadataBase(address).plus(OFFSET_VEC_OFFSET);
  }
  
  public static Address getFirstObjVecBase(Address address) {
    // Comes exactly after OffsetVec
    return getMetadataBase(address).plus(FIRSTOBJ_VEC_OFFSET);
  }
  
  public static Address getStatusVecBase(Address address) {
    return getMetadataBase(address).plus(STATUS_VEC_OFFSET);
  }

  /**
   * Given an address, return the address of the live word for
   * that address.
   *
   * @param address The address whose live word address is to be returned
   * @return The address of the live word for this object
   */
  @Inline
  public static Address getLiveWordAddress(Address address) {
    Extent liveWordOffset = EmbeddedMetaData.getMetaDataOffset(
        address, LOG_LIVE_COVERAGE/* - LOG_BYTES_IN_WORD*/, LOG_BYTES_IN_WORD);
    return getLiveMapBase(address).plus(liveWordOffset);
  }
  @Inline
  public static Address getOffsetVectorCellAddress(Address address) {
    Extent cellOffset = EmbeddedMetaData.getMetaDataOffset(
        address, LOG_OFFSET_COVERAGE/* - LOG_BYTES_IN_ADDRESS*/, LOG_BYTES_IN_ADDRESS);
    return getOffsetVecBase(address).plus(cellOffset);
  }

  public static Address getFirstObjCellAddress(Address address) {
    Extent cellOffset = EmbeddedMetaData.getMetaDataOffset(
        address, LOG_FIRSTOBJ_COVERAGE/* - LOG_BYTES_IN_ADDRESS*/, LOG_BYTES_IN_ADDRESS);
    return getFirstObjVecBase(address).plus(cellOffset);
  }

  public static Address getStatusVectorCellAddress(Address address) {
    Extent cellOffset = EmbeddedMetaData.getMetaDataOffset(
        address, LOG_STATUS_COVERAGE/* - LOG_BYTES_IN_ADDRESS*/, LOG_BYTES_IN_ADDRESS);
    return getStatusVecBase(address).plus(cellOffset);
  }



}